import torch.nn as nn
from torchvision.models.mobilenet import mobilenet_v2

class DigitsMobilenet(nn.Module):
    def __init__(self, class_num=11):
        super(DigitsMobilenet, self).__init__()

        self.net = mobilenet_v2(pretrained=True).features
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.bn = nn.BatchNorm1d(1280)
        self.fc1 = nn.Linear(1280, class_num)
        self.fc2 = nn.Linear(1280, class_num)
        self.fc3 = nn.Linear(1280, class_num)
        self.fc4 = nn.Linear(1280, class_num)

    def forward(self, img):
        features = self.avgpool(self.net(img)).view(-1, 1280)
        features = self.bn(features)

        fc1 = self.fc1(features)
        fc2 = self.fc2(features)
        fc3 = self.fc3(features)
        fc4 = self.fc4(features)

        return fc1, fc2, fc3, fc4 